import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from tqdm import tqdm

from config import parser
from datalist import VideoDataset2
from model import C3D


best_acc = 0


class predict(object):
    def __init__(self):
        self.args = parser.parse_args()
        print(f"-----------{self.args.project_name}-------------")

        use_cuda = self.args.use_cuda and torch.cuda.is_available()
        if use_cuda:
            torch.cuda.manual_seed(self.args.seed)  # 为当前GPU设置随机种子
        else:
            torch.manual_seed(self.args.seed)  # 为CPU设置种子用于生成随机数，以使得结果是确定的
        self.device = torch.device("cuda" if use_cuda else "cpu")
        train_kwargs = {'num_workers': 0, 'pin_memory': True} if use_cuda else {}
        test_kwargs = {'num_workers': 0, 'pin_memory': False} if use_cuda else {}
        '''
        构造DataLoader
        '''
        self.train_dataloader = DataLoader(VideoDataset2(self.args.train_dir), batch_size=self.args.train_batch_size,
                                           shuffle=True, **train_kwargs)
        self.test_dataloader = DataLoader(VideoDataset2(self.args.test_dir), batch_size=self.args.test_batch_size,
                                          shuffle=True, **test_kwargs)

        self.model_rgb = C3D(2).to(self.device)
        self.model_binary = C3D(2).to(self.device)
        self.model_flow = C3D(2).to(self.device)

        if use_cuda:
            self.model_rgb = torch.nn.DataParallel(self.model_rgb, device_ids=range(torch.cuda.device_count()))
            self.model_flow = torch.nn.DataParallel(self.model_flow, device_ids=range(torch.cuda.device_count()))
            self.model_binary = torch.nn.DataParallel(self.model_binary, device_ids=range(torch.cuda.device_count()))
            cudnn.benchmark = True

        self.model_rgb.load_state_dict(torch.load("./weights/Epoch-260-best_acc-0.9350000023841858_loss_0.23.pth", map_location=self.device)['model_state_dict'])
        self.model_binary.load_state_dict(
            torch.load("./weights_binary/Epoch-120-best_acc-0.925000011920929_loss_0.27.pth", map_location=self.device)['model_state_dict'])
        self.model_flow.load_state_dict(
            torch.load("./weights_flow/Epoch-326-best_acc-0.7300000190734863_loss_0.69.pth", map_location=self.device)['model_state_dict'])


        self.predict(epoch=1)

        torch.cuda.empty_cache()
        print("finish model testing")


    def predict(self, epoch):

        global best_acc

        self.model_rgb.eval()
        self.model_flow.eval()
        self.model_binary.eval()

        correct = torch.zeros(1).squeeze().cuda()
        total = torch.zeros(1).squeeze().cuda()

        average_loss = []

        pbar = tqdm(self.test_dataloader,
                    desc=f'Test Epoch{epoch}/{self.args.epoches}',
                    mininterval=0.3)
        for data, target in pbar:
            data, target = data.to(self.device), target.to(self.device)
            with torch.no_grad():
                output_rgb = self.model_rgb(data)
                output_binary = self.model_binary(data)
                output_flow = self.model_flow(data)

                output=output_rgb+output_binary+output_flow

                pred = torch.argmax(output, 1)

                correct += (pred == target).sum().float()
                total += len(target)
                predict_acc = correct / total
            pbar.set_description(
                f'Test Epoch:{epoch}/{self.args.epoches} acc:{predict_acc}')



predict()
